#ifndef _MEMORY_H_
#define _MEMORY_H_

#include <iostream>
#include <fstream>
#include <string.h>
#include <memory>
#include <sys/ipc.h>
#include <sys/shm.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>
#include <list>
#include <atomic>
#include <Python.h>
#include <opencv2/opencv.hpp>
#include "data.h"

namespace sp{

// type
static const uint8_t Data_Memory = 0x00;
static const uint8_t Cv_Mat_Data_3840_2160_Memory = 0x01;
static const uint8_t Cv_Mat_Data_1920_1080_Memory = 0x02;
static const uint8_t Cv_Mat_Data_1280_720_Memory = 0x03;
static const uint8_t Cv_Mat_Data_960_540_Memory = 0x04;
static const uint8_t Tensor_Data_1K_Memory = 0x05;
static const uint8_t Tensor_Data_10K_Memory = 0x06;
static const uint8_t Tensor_Data_100K_Memory = 0x07;
static const uint8_t Tensor_Data_1M_Memory = 0x08;
static const uint8_t Tensor_Data_10M_Memory = 0x09;
static const uint8_t Tensor_Data_100M_Memory = 0x0a;
static const uint8_t Tensor_Data_500M_Memory = 0x0b;

/**
    共享内存中的基本类型数组
*/
template<typename T>
class SharedMemory{
private:
    std::string name;           // 名称
    size_t element_size;        // 存储元素的大小
    long capacity;              // 容量
    key_t shm_id;               // 共享内存句柄
    char* ptr;                  // 起始地址
    std::atomic<bool>* occ_ptr; // occ起始地址
    T* data_ptr;                // 数据起始地址
public:
    SharedMemory(std::string name, long capacity);
    ~SharedMemory();
    void setValue(long pos, T value);
    T getValue(long pos);
};

// 存储统一种类型的共享内存池
/* 内存布局
    | occ | element0 | element1 | ... | elementn |
*/

class SharedMemoryPool{
private:
    uint8_t type;               // 存储元素类型
    std::string name;           // 名称
    size_t element_size;        // 存储元素大小
    long capacity;              // 容量
    key_t shm_id;               // 共享内存句柄
    char* ptr;                  // 起始地址
    std::atomic<bool>* occ_ptr; // occ起始地址
    char* data_ptr;             // 数据起始地址
    std::ofstream* fout_ptr;    // 文件输出
public:
    SharedMemoryPool(uint8_t type, std::string name, long capacity);
    SharedMemoryPool(uint8_t type, std::string name, size_t element_size, long capacity);
    ~SharedMemoryPool();
    void bindFileOut(std::ofstream& fout);
    size_t typeToElementSize(uint8_t type);
    uint8_t getType();
    char* getDataPtr();
    long getCapacity();
    bool contain(char* addr);
    char* malloc();
    void free(char* addr);
    long getFreeCount();
};

class SharedMemoryManager{
private:
    // 共享内存池集合
    std::map<std::string, SharedMemoryPool*> n_pools_map;  // 普通用途
    std::map<std::string, SharedMemoryPool*> o_pools_map;  // 固定用途
    // 文件输出
    std::ofstream fout;
    // 私有方法
    uint8_t addrToType(void* addr);
public:
    SharedMemoryManager();
    ~SharedMemoryManager();
    size_t typeToElementSize(uint8_t type);
    SharedMemoryPool createSharedMemoryPool(std::string name, uint8_t type, long cnt);
    SharedMemoryPool createSharedMemoryPool(std::string name, uint8_t type, size_t element_size, long cnt);
    char* malloc(uint8_t type);
    char* malloc(size_t size);
    void free(char* ptr);             // 先减引用计数，减为0再释放
    void check();                     // 输出共享内存的状态
    void writeToFile(std::string move, std::string name, void* addr);    // 记录到文件
    static std::string typeToString(uint8_t type);
};

// DONE: 完善GPU显存管理
class Gpu_SharedMemoryManager{
private:
    // 显存划分相关变量，显存布局 | DATA | MODEL | COMPUTION |
    size_t TOTAL_SHARED_MEMORY_SIZE;     // 总显存大小
    size_t DATA_SHARED_MEMORY_SIZE;      // DATA区显存大小
    int DATA_STREAM_COUNT;               // DATA区划分stream数量
    size_t MAX_DATA_STREAM_COUNT;        // DATA区划分stream数量最大值
    size_t DATA_EVERY_STREAM_SIZE;       // DATA区每个stream的大小
    size_t MODEL_SHARED_MEMORY_SIZE;     // MODEL区显存大小
    size_t COMPUTION_SHARED_MEMORY_SIZE; // COMPUTION区显存大小
    int GPU_NUMBER;                      // GPU个数
    // shm标识符
    key_t stream_occ_shm_id;
    // 对象私有变量
    char* stream_occ_shm_ptr;            // 表示stream占用信息的内存首地址
    std::atomic<bool>** stream_occ_ptr;  // 表示stream占用信息的atomic指针
    // block map
    std::unordered_map<short, std::pair<size_t, size_t>> blocks_map; // 存储每个block的id到pos, size对的映射，其中block 0,...,n代表params blocks, -1 代表calculate block
public:
    Gpu_SharedMemoryManager();
    ~Gpu_SharedMemoryManager();
    void cuda_env_init(short gpu_id);           // 初始化cuda环境
    void allotSharedCache();                    // 申请SharedCache
    void freeSharedCache();                     // 释放SharedCache
    void sendSharedCache();                     // 发送SharedCache句柄
    void recvSharedCache();                     // 接受SharedCache句柄
    short getStreamId(size_t pos);              // 获取streamId。
    size_t getPos(short stream_id);             // 获取stream显存起始地址。
    size_t getSize(short stream_id);            // 获取stream大小。
    short malloc(short gpu_id);                 // 申请一块空间，非阻塞，有则返回stream_id，没有则返回-1
    void free(short gpu_id, short stream_id);   // 释放一块空间
};

};

#endif
